import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy import stats
from statsmodels.tsa.arima.model import ARIMA
from sklearn.metrics import mean_squared_error
import matplotlib as mpl
from statsmodels.tsa.stattools import adfuller as ADF
mpl.use('TkAgg')

# ---------------数据预处理------------------------------
df = pd.read_csv('data/factor.csv', dtype=str, index_col=0)
aqi = pd.read_csv('data/aqi.csv', index_col=0)

df['date'] = pd.to_datetime(df['Year'] + ' ' + df['Month'] + ' ' + df['Day'], format='%Y %m %d')
df.drop(['Year', 'Month', 'Day'], axis=1, inplace=True)

df1 = df.iloc[:, -1:]
df2 = df.iloc[:, :-1]

df = pd.concat([df1, df2, aqi], axis=1)
df.to_csv("data/Q3data.csv")

# ---------------读取数据------------------------------
df = pd.read_csv("data/Q2.csv", index_col=0)
# 将df序号设置为date
df['AQI'] = df['AQI'].interpolate()
df.index = df['date']
del df['date']

df = df.iloc[:, -1:]
print(df)

df.plot(figsize=(16, 8))
plt.savefig("img/AQI-time.png")
# plt.show()

# 1、一阶差分
df_diff = df.diff(1)
df_diff = df_diff.dropna()

# 绘制 ACF 和 PACF图像
fig = plt.figure(figsize=(12, 8))
ax1 = fig.add_subplot(211)
fig = sm.graphics.tsa.plot_acf(df_diff.values.squeeze(), lags=40, ax=ax1)
ax2 = fig.add_subplot(212)
fig = sm.graphics.tsa.plot_pacf(df_diff, lags=40, ax=ax2)
plt.savefig("img/ACF-PCF.png")

# 2、数据平稳性检验
before = ADF(df["AQI"])
print(before)
# 返回值依次为adf、pvalue、usedlag、nobs、critical  values、icbest、regresults、resstore  p<0.05时表示稳定
# (-7.079692912376129, 4.701352703395662e-10, 28, 3012, {'1%': -3.43252293552353, '5%': -2.862500063122179, '10%': -2.5672810666012356}, 29239.301366837124)
after_diff = ADF(df_diff["AQI"])
print(after_diff)
# (-15.245591941093485, 5.091908567485942e-28, 26, 3013, {'1%': -3.432522213721007, '5%': -2.862499744326095, '10%': -2.567280896877994}, 29277.287511007988)
# 原始数据平稳性检验pvalue < 0.05 是平稳序列，故d = 0

# 3、根据bic/aic 定p、q

pmax = 10  # 一般阶数不超过length/10
qmax = 10  # 一般阶数不超过length/10

bic_matrix = []  # bic矩阵
for p in range(pmax + 1):
    tmp = []
    for q in range(qmax + 1):  # 存在部分报错，所以用try来跳过报错。
        try:
            tmp.append(ARIMA(df["AQI"], order=(p, 0, q)).fit().bic)
        except:
            tmp.append(None)
    bic_matrix.append(tmp)
bic_matrix = pd.DataFrame(bic_matrix)  # 从中可以找出最小值
p, q = bic_matrix.stack().idxmin()
# #先用stack展平，然后用idxmin找出最小值位置。
print(u'BIC最小的p值和q值为：%s、%s' % (p, q))
# 得到最佳参数为201

pmax = 10
qmax = 10
aic_matrix = []  # aic矩阵
for p in range(pmax + 1):
    tmp = []
    for q in range(qmax + 1):  # 存在部分报错，所以用try来跳过报错。
        try:
            tmp.append(ARIMA(df["AQI"], order=(p, 0, q)).fit().aic)
        except:
            tmp.append(None)
    aic_matrix.append(tmp)
aic_matrix = pd.DataFrame(bic_matrix)  # 从中可以找出最小值
p, q = bic_matrix.stack().idxmin()
# #先用stack展平，然后用idxmin找出最小值位置。
print(u'AIC最小的p值和q值为：%s、%s' % (p, q))
# 找到最佳参数为2 0 1

# 4、模型训练
# 定义一个空列表，用于存储预测值

window = 16
train_proportion = 0.8


n_train = round(df.shape[0] * train_proportion)
train = df.iloc[:n_train, :]    # 2433
test = df.iloc[n_train:, :]     # 608 4*152  16*38

prediction = list()
obs_list = list()
history = pd.DataFrame()
pred = pd.DataFrame()
# 定义一个循环，对每个时间窗口内的数据，创建并拟合 ARIMA 模型，并预测下一个时间点的值
for i in range(38):
    history = df[(window * i):(window*i + n_train)] # 取出当前时间窗口内的数据
    for t in range(window):
        model = ARIMA(history, order=(2, 0, 1)) # 创建 ARIMA(2,0,1) 模型
        model_fit = model.fit() # 拟合模型

        output = model_fit.forecast() # 预测下一个时间点的值
        yhat = output[0] # 取出预测值
        prediction.append(yhat) # 将预测值添加到列表中

        test_list = test['AQI'].values
        obs = test_list[i*window+t] # 取出实际值
        obs_list.append(obs)

        temp_obs = {"AQI": obs}
        temp_pred = {"AQI": yhat}
        index = test.index[i*window+t]
        df_pred = pd.DataFrame(temp_pred, index=[index])
        df_temp = pd.DataFrame(temp_obs, index=[index])
        history = pd.concat([history, df_temp], axis=0)
        pred = pd.concat([pred, df_pred], axis=0)
    if i == 37:
        pred_ = model.fit().forecast(steps=12)
# history.to_csv("history.csv")
pred.to_csv("pred.csv")




